import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dense,SimpleRNN,Embedding
import os


input_word='abcde'
w_to_id={'a':0,'b':1,'c':2,'d':3,'e':4}
# id_to_onehot={0:[1,0,0.,0,0],1:[0,1,0,0.,0],2:[0,0,1,0.,0],3:[0,0,0.,1,0],4:[0,0,0.,0,1]}

# x_train=[
#     id_to_onehot[w_to_id['a']],
#     id_to_onehot[w_to_id['b']],
#     id_to_onehot[w_to_id['c']],
#     id_to_onehot[w_to_id['d']],
#     id_to_onehot[w_to_id['e']]
#     ]
x_train=[
    w_to_id['a'],
    w_to_id['b'],
    w_to_id['c'],
    w_to_id['d'],
    w_to_id['e'],
]

y_train=[
    w_to_id['b'],
    w_to_id['c'],
    w_to_id['d'],
    w_to_id['e'],
    w_to_id['a'],
]

np.random.seed(7)
np.random.shuffle(x_train)
np.random.seed(7)
np.random.shuffle(y_train)
tf.random.set_seed(7)

# Embedding输入要求[送入样本数,循环核时间展开步数]
x_train=np.reshape(x_train, (len(x_train),1))
y_train=np.array(y_train)

model=tf.keras.Sequential([
    Embedding(5,2),
    SimpleRNN(3),
    Dense(5,activation='softmax')
])
model.compile(optimizer=tf.keras.optimizers.Adam(0.01),
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
              metrics=['sparse_categorical_accuracy'])

# print(os.path.dirname(os.path.abspath("checkpoint/rnn_onhot.ckpt")))
checkpoint_save_path=os.path.abspath("checkpoint/rnn_embedding_1pre1.ckpt")
print(checkpoint_save_path)
if os.path.exists(checkpoint_save_path+'.index'):
    print('-------------------load the model-------------------')
    model.load_weights(checkpoint_save_path)

cp_callback=tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_save_path,
                                                save_weights_only=True,
                                                save_best_only=True,
                                                monitor='loss')       

history=model.fit(x_train,y_train,batch_size=32,epochs=100,callbacks=[cp_callback])
model.summary()

file=open('./weights_embedding_1pre1.txt','w')
for v in model.trainable_variables:
    file.write(str(v.name)+'\n')
    file.write(str(v.shape)+'\n')
    file.write(str(v.numpy())+'\n')
file.close()



preNum=int(input("输入测试字母数量："))
for i in range(preNum):
    alphabet1=input("输入测试字母：")
    alphabet=[w_to_id[alphabet1]]
    alphabet=np.reshape(alphabet, [1,1])
    result=model.predict(alphabet)
    pred=tf.argmax(result,axis=1)
    pred=int(pred)
    tf.print(alphabet1+'->'+input_word[pred])